package com.bigdata.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, IntegerType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

object SparkSQL_Avg_Age_UDAF_3 {

    def main(args: Array[String]): Unit = {

        // TODO 计算平均年龄
        val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("sparkSQL")
        val spark: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()

        val df: DataFrame = spark.read.json("datas/user.json")
        df.createOrReplaceTempView("user")
        val myAvgUDAF = new MyAvgUDAF
        spark.udf.register("ageAvg", myAvgUDAF)
        spark.sql("select ageAvg(age) from user").show

        spark.stop()
    }

    class MyAvgUDAF extends UserDefinedAggregateFunction{

        // 聚合函数输入参数的数据类型
        override def inputSchema: StructType = {
            StructType(
                Array(
                    StructField("age", IntegerType)
                )
            )
        }

        // 聚合函数缓冲区中值的数据类型(age,count)
        override def bufferSchema: StructType = {
            StructType(
                Array(
                    StructField("sum", LongType),
                    StructField("count", LongType)
                )
            )
        }

        // 函数返回值的数据类型
        override def dataType: DataType = DoubleType

        // 稳定性:对于相同的输入是否一直返回相同的输出。
        override def deterministic: Boolean = true


        // 函数缓冲区初始化
        override def initialize(buffer: MutableAggregationBuffer): Unit = {
            // 存年龄的总和
            buffer(0) = 0L

            // 存年龄的个数
            buffer(1) = 0L
        }

        // 更新缓冲区中的数据
        override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
            buffer(0) = buffer.getLong(0) + input.getInt(0)
            buffer(1) = buffer.getLong(1) + 1
        }

        // 合并缓冲区
        override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
            buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
            buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
        }

        // 计算最终结果
        override def evaluate(buffer: Row): Any = {
            buffer.getLong(0).toDouble / buffer.getLong(1)
        }
    }
}
